# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2024/12/4 14:33
# @Function: 分析land30m的数据tif图的工具
import os
import rasterio
import numpy as np
from collections import Counter
import matplotlib
from matplotlib import pyplot as plt
from PIL import Image


# 调用函数并打印结果
tif_path = r'E:\land30m2010\beijing\merged\beijing_land30m.tif'  # 替换为你的TIF文件路径

# 定义一级分类的值范围
category_mapping = {
    1: [101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112],  # 林地
    2: [21, 22, 23, 24],  # 草地
    3: [31, 32, 33, 34, 35, 36, 37],  # 湿地
    4: [41, 42],  # 耕地
    5: [51, 52, 53, 54],  # 人工表面
    6: [61, 62, 63, 64, 65, 66, 67, 68, 69],  # 其它
}


# 读取TIF文件
def read_tif(tif_path):
    with rasterio.open(tif_path) as src:
        return src.read(1)  # 读取第一波段数据


# 获取一级分类
def get_level_1_category(value):
    for level_1, values in category_mapping.items():
        if value in values:
            return level_1
    return 0  # 如果没有对应的一级分类，返回0表示无效分类


# 统计每个一级分类的像素数量及其占比
def get_category_stats(tif_path):
    data = read_tif(tif_path)

    # 忽略无效值（255），将其替换为0
    valid_data = np.where(data != 255, data, 0)

    # 将所有像素值映射为一级分类
    level_1_categories = np.vectorize(get_level_1_category)(valid_data)

    # 统计每个一级分类的数量
    counter = Counter(level_1_categories.flatten())

    # 计算总有效像素数（排除255的像素）
    total_valid_pixels = (valid_data != 0).sum()

    # 计算各一级分类的占比
    category_stats = {}
    for category, count in counter.items():
        if category != 0:  # 忽略无效类别
            category_stats[category] = {
                'count': count,
                'percentage': (count / total_valid_pixels) * 100
            }

    return category_stats


# 配置matplotlib使用支持中文的字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体
matplotlib.rcParams['axes.unicode_minus'] = False   # 防止负号显示为乱码
# 绘制饼图
def plot_pie_chart(category_stats):
    categories = list(category_stats.keys())
    percentages = [category_stats[category]['percentage'] for category in categories]

    # 绘制饼图
    plt.figure(figsize=(8, 8))
    plt.pie(percentages, labels=categories, autopct='%1.1f%%', startangle=140, colors=plt.cm.Paired.colors)
    plt.title('土地覆被分类占比')
    plt.axis('equal')  # 保证饼图是圆形的
    plt.show()

category_stats = get_category_stats(tif_path)

# 打印结果
for category, stats in category_stats.items():
    print(f"一级分类 {category} 的像素数量: {stats['count']}, 占比: {stats['percentage']:.2f}%")

# 绘制饼图
# plot_pie_chart(category_stats)


'''
=====================================类别空间分布图================================
'''

# 1. 一级分类到颜色的映射
category_colors = {
    1: (0, 128, 0),  # 林地 -> 深绿色
    2: (0, 255, 0),  # 草地 -> 绿色
    3: (0, 0, 255),  # 湿地 -> 蓝色
    4: (255, 255, 0),  # 耕地 -> 黄色
    5: (255, 165, 0),  # 人工表面 -> 橙色
    6: (128, 0, 128),  # 其它 -> 紫色
}

# 3. 根据二级分类映射转换为一级分类
def get_primary_category(sub_category):
    for primary, sub_categories in category_mapping.items():
        if sub_category in sub_categories:
            return primary
    return 0  # 默认返回0，表示未知类别


# 4. 将TIF图像转换为指定颜色的PNG图像
def convert_tif_to_colored_png(tif_path, png_path):
    # 读取TIF图像
    img = Image.open(tif_path)

    # 将图像转换为灰度模式（假设原图每个像素值表示一个类别）
    img_array = np.array(img)

    # 创建一个空的RGBA图像，第四个通道用于透明度
    colored_img_array = np.zeros((img_array.shape[0], img_array.shape[1], 4), dtype=np.uint8)

    # 遍历每个像素，将二级分类值转换为一级分类，然后映射到对应的颜色
    for i in range(img_array.shape[0]):
        for j in range(img_array.shape[1]):
            sub_category = img_array[i, j]
            primary_category = get_primary_category(sub_category)
            if primary_category != 0:  # 如果不是未知类别
                colored_img_array[i, j, :3] = category_colors.get(primary_category, (0, 0, 0))  # RGB
                colored_img_array[i, j, 3] = 255  # 完全不透明
            else:
                colored_img_array[i, j] = [0, 0, 0, 0]  # 透明色

    # 创建一个PIL图像对象
    colored_img = Image.fromarray(colored_img_array, 'RGBA')

    # 保存为PNG格式
    colored_img.save(png_path, 'PNG')
    print(f"图像已保存为 {png_path}")

file_name, _ = os.path.splitext(os.path.basename(tif_path))  # 分离文件名和扩展名
# 示例调用
# convert_tif_to_colored_png(tif_path, f'./tif_out/land30m_{file_name}.png')